# -*- coding: utf-8 -*-
# @Time    : 2020/12/4 9:38
# @Author  : DuJiabao
# @File    : PhotoCombiner.py
# @Description : 这是一个图像合成的文件。
import numpy as np
import cv2
from ImageProcess import picRotate, cropPic, synthesis
import xml.etree.ElementTree as ET
import os


class Blackground(object):
    def __init__(self, blackground: np.ndarray):
        self.blackground = blackground
        self.height, self.width, _ = self.blackground.shape

    def getHW(self):
        """
        获取当前背景的高度、宽度
        :return:
        """
        return self.height, self.width

    def getHeight(self):
        return self.height

    def getWidth(self):
        return self.width

    def getBlackground(self):
        return self.blackground.copy()

    def rotate(self, angle):
        self.blackground = picRotate(self.blackground, angle,
                                     tuple(self.blackground[int(self.height / 8), int(self.width / 8)].tolist()))
        self.height, self.width, _ = self.blackground.shape

    def show(self):
        cv2.imshow("show", self.blackground)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def setBlackground(self, blackground):
        self.blackground = blackground
        self.height, self.width, _ = self.blackground.shape

    def resize(self, size: tuple):
        """
        resize blackground
        :param size: 新的size，(height,width)
        :return:
        """
        self.blackground = cv2.resize(self.blackground, size)
        self.height, self.width, _ = size


class BlackgroundFromPath(Blackground):
    def __init__(self, path):
        blackground = cv2.imread(path)
        super().__init__(blackground)


class BoundingBox(object):
    def __init__(self, xmin: int, ymin: int, xmax: int, ymax: int):
        """
        初始化BoundingBox
        :param xmin:
        :param ymin:
        :param xmax:
        :param ymax:
        """
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax

    def set(self, xmin: int, ymin: int, xmax: int, ymax: int):
        """
        重新设置BoundingBox
        :param xmin:
        :param ymin:
        :param xmax:
        :param ymax:
        :return:
        """
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax

    def getXYXY(self) -> tuple:
        """
        依照XYXY顺序返回
        :return: xmin,ymin,xmax,ymax
        """
        return self.xmin, self.ymin, self.xmax, self.ymax

    def getYXYX(self) -> tuple:
        """
        依照YXYX顺序返回
        :return: ymin,xmin,ymax,xmax
        """
        return self.xmin, self.ymin, self.xmax, self.ymax

    def getXXYY(self) -> tuple:
        """
        依照XXYY顺序返回
        :return: xmin,xmax,ymin,ymax
        """
        return self.xmin, self.xmax, self.ymin, self.ymax

    def getYYXX(self) -> tuple:
        """
        依照YYXX顺序返回
        :return: ymin,ymax,xmin,xmax
        """
        return self.ymin, self.ymax, self.xmin, self.xmax

    def getXY_XY(self) -> tuple:
        """
        依照XY_XY的顺序返回
        :return: (xmin,ymin),(xmax,ymax)
        """
        return (self.xmin, self.ymin), (self.xmax, self.ymax)

    def getHW(self) -> tuple:
        return self.ymax - self.ymin, self.xmax - self.xmin

    def getHeight(self):
        return self.ymax - self.ymin

    def getWidth(self):
        return self.xmax - self.xmin

    def getCenter(self) -> tuple:
        """
        返回中心
        :return: y_center, x_center
        """
        return (self.ymax - self.ymin) / 2, (self.xmax - self.xmin) / 2

    def displace(self, x: int, y: int):
        """
        实现bounding
        :param y: y方向上的偏移
        :param x: x方向上的偏移
        :return:
        """
        self.xmin += x
        self.xmax += x
        self.ymin += y
        self.ymax += y


class Item(object):
    def __init__(self, item: np.ndarray, cls: str):
        """
        初始化用于被粘贴的item。空白处的值为0。
        :param item: item的图片
        :param cls: 类别
        :param isFull: item是否完全覆盖图片
        """
        self.item = item
        self.cls = cls
        self.height, self.width, _ = item.shape
        self.bbox = self.findBoundingBox(item)
        self.deleteBlank()

    # 可实现在外部调用，但无需实例化
    @staticmethod
    def findBoundingBox(img: np.ndarray):
        """
        寻找图片的最小外接矩形，空白处为0
        :param img:
        :return: BoundingBox
        """
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # 二值化
        ret, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
        contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) != 1:
            print("Contours != 1")
        x, y, w, h = cv2.boundingRect(contours[0])
        return BoundingBox(x, y, x + w, y + h)

    def rotate(self, angle):
        """
        实现item的随机旋转
        :param angle: 旋转的角度，顺时针
        :return:
        """
        self.item = picRotate(self.item, angle, (0, 0, 0))
        self.bbox = self.findBoundingBox(self.item)
        self.deleteBlank()

    def show(self):
        """
        显示目前图片状态
        :return:
        """
        temp = self.item.copy()
        temp = cv2.rectangle(temp, *self.bbox.getXY_XY(), (255, 0, 0))
        cv2.imshow("show", temp)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def deleteBlank(self):
        """
        删除图片的空白区域
        :return:
        """
        self.item = cropPic(self.item, *self.bbox.getXYXY())
        self.height, self.width, _ = self.item.shape
        self.bbox.set(0, 0, self.width, self.height)

    def getHW(self):
        return self.height, self.width

    def getHeight(self):
        return self.height

    def getWidth(self):
        return self.width

    def getItem(self):
        return self.item.copy()

    def resize(self, size: tuple):
        """
        resize item
        :param size: 新的size，(height,width)
        :return:
        """
        self.item = cv2.resize(self.item, size)
        self.height, self.width, _ = size
        self.bbox.set(0, 0, self.width, self.height)


class ItemFromPath(Item):
    def __init__(self, path, cls: str):
        item = cv2.imread(path)
        super().__init__(item, cls)


class Object(BoundingBox):
    def __init__(self, cls: str, xmin: int, ymin: int, xmax: int, ymax: int):
        super(Object, self).__init__(xmin, ymin, xmax, ymax)
        """
        用于存储目标的类别和位置信息
        :param cls: 类别
        :param xmin:
        :param ymin:
        :param xmax:
        :param ymax:
        """
        self.cls = cls

    def getCls(self):
        return self.cls


class OnePic(object):
    def __init__(self, blackground: Blackground):
        """
        包含图片、目标及其标注的类
        :param blackground: 背景
        """
        self.blackground = blackground
        self.objects = []

    def loadObjectFromXML(self, path: str):
        """
        从xml文件导入标注信息
        :param path: 存放xml文件的路径
        :return:
        """
        in_file = open(path)
        tree = ET.parse(in_file)
        root = tree.getroot()
        for obj in root.iter('object'):
            cls = obj.find('name').text
            xmlbox = obj.find('bndbox')
            xmin = int(xmlbox.find('xmin').text)
            ymin = int(xmlbox.find('ymin').text)
            xmax = int(xmlbox.find('xmax').text)
            ymax = int(xmlbox.find('ymax').text)
            self.objects.append(Object(cls, xmin, ymin, xmax, ymax))

    def show(self):
        temp = self.blackground.getBlackground()
        # 线粗
        tl = round(0.002 * (temp.shape[0] + temp.shape[1]) / 2) + 1
        tf = max(tl - 1, 1)  # font thickness

        for oj in self.objects:
            cv2.rectangle(temp, *oj.getXY_XY(), (0, 0, 255), thickness=tl, lineType=cv2.LINE_AA)
            cv2.putText(temp, oj.cls, (oj.xmin, oj.ymin - 2), 0, tl / 3, [225, 255, 255], thickness=tf,
                        lineType=cv2.LINE_AA)
        cv2.namedWindow("show", cv2.WINDOW_FREERATIO)
        cv2.imshow("show", temp)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def addItem(self, item: [Item, ItemFromPath], center: tuple, size=None, angle=None):
        """
        在背景上添加item
        :param item: Item/ItemFromPath类，在背景上粘贴的内容
        :param center: item的中心点坐标，(y, x)
        :param size: item resize 之后的shape，(height,width)，默认为None
        :param angle: 旋转的角度，顺时针为正，int，默认为None
        :return:
        """
        if angle is not None:
            item.rotate(angle)
        if size is not None:
            item.resize(size)

        if item.getHeight() > self.blackground.getHeight() or item.getWidth() > self.blackground.getWidth():
            print("The item is too big!")
        # 粘贴
        bg, bbox = synthesis(self.blackground.getBlackground(), item.getItem(), center)
        self.blackground.setBlackground(bg)  # 重设bg
        ymin, ymax, xmin, xmax = bbox
        # 滤去被完全遮挡的目标
        objects = []
        for oj in self.objects:
            xmin_0, ymin_0, xmax_0, ymax_0 = oj.getXYXY()
            if xmin <= xmin_0 and ymin <= ymin_0 and xmax >= xmax_0 and ymax >= ymax_0:
                continue
            objects.append(oj)
        # 加入新目标位置
        objects.append(Object(item.cls, xmin, ymin, xmax, ymax))
        self.objects = objects

    def savePic(self, path: str):
        cv2.imwrite(path, self.blackground.getBlackground())

    def saveXML(self, path: str):
        filename = os.path.split(path)[1].replace("xml", "jpg")
        root = ET.Element('annotation')

        folder_root = ET.SubElement(root, "folder")
        folder_root.text = "images"

        filename_root = ET.SubElement(root, "filename")
        filename_root.text = filename

        path_root = ET.SubElement(root, "path")
        path_root.text = os.path.abspath(path).replace("Annotations", "images").replace("xml", "jpg")

        source_root = ET.SubElement(root, "source")
        database_root = ET.SubElement(source_root, "database")
        database_root.text = "Unknown"

        size_root = ET.SubElement(root, "size")
        height_root = ET.SubElement(size_root, "height")
        height_root.text = str(self.blackground.getHeight())
        width_root = ET.SubElement(size_root, "width")
        width_root.text = str(self.blackground.getWidth())
        depth_root = ET.SubElement(size_root, "depth")
        depth_root.text = "3"

        for oj in self.objects:
            ymin, ymax, xmin, xmax = oj.getYYXX()
            object_root = ET.SubElement(root, "object")
            name_root = ET.SubElement(object_root, "name")
            name_root.text = oj.getCls()
            pose_root = ET.SubElement(object_root, "pose")
            pose_root.text = "Unspecified"
            truncated_root = ET.SubElement(object_root, "truncated")
            truncated_root.text = "0"
            difficult_root = ET.SubElement(object_root, "difficult")
            difficult_root.text = "0"
            bndbox_root = ET.SubElement(object_root, "bndbox")
            xmin_root = ET.SubElement(bndbox_root, "xmin")
            ymin_root = ET.SubElement(bndbox_root, "ymin")
            xmax_root = ET.SubElement(bndbox_root, "xmax")
            ymax_root = ET.SubElement(bndbox_root, "ymax")
            ymin_root.text = str(ymin)
            ymax_root.text = str(ymax)
            xmin_root.text = str(xmin)
            xmax_root.text = str(xmax)
        tree = ET.ElementTree(root)

        tree.write(path, encoding='utf-8')

    def saveTxt(self, path: str, clsList: list):
        """
        在txt文件中保存标注信息
        :param path: 保存地址
        :param clsList: 类别目录
        :return:
        """
        with open(path, "w") as f:
            for oj in self.objects:
                # 需要将bbox的坐标归一化
                f.write(f"{clsList.index(oj.getCls())} "
                        f"{oj.getCenter()[1] / self.blackground.getWidth()} "
                        f"{oj.getCenter()[0] / self.blackground.getHeight()} "
                        f"{oj.getWidth() / self.blackground.getWidth()} "
                        f"{oj.getHeight() / self.blackground.getHeight()}\n")

    def saveWithXML(self, folder, name):
        self.savePic(os.path.join(folder, "images", name + ".jpg"))
        self.saveXML(os.path.join(folder, "Annotations", name + ".xml"))
        self.saveTxt(os.path.join(folder, "labels", name + ".txt"), ["corn", "pot"])


if __name__ == '__main__':
    pic = OnePic(BlackgroundFromPath("./dataset/images/WIN_20201109_14_44_07_Pro.jpg"))
    pic.loadObjectFromXML("./dataset/images/WIN_20201109_14_44_07_Pro.xml")
    pot = ItemFromPath("./dataset/images/8.png", "pot")
    pic.addItem(pot, (1000, 1000))
    # pic.show()
    pic.saveWithXML("./dataset/", "1")

    # img = cv2.imread("./images/8.png")
    # bg = Blackground(img)
    # bg.show()
    # bg.rotate(30)
    # bg.show()
    # pot = Item(img, "pot")
    # # pot.rotate(100)
    # pot.show()
    # cv2.rectangle(img, (pot.bbox.xmin, pot.bbox.ymin), (pot.bbox.xmax, pot.bbox.ymax), (255, 0, 0))
    # cv2.imshow("img", img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
